package cn.tedu.demos.webSocket;

import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@ServerEndpoint("/websocket-endpoint")
@Component
public class CustomWebSocketEndpoint {

    // 存储所有连接的 WebSocket 会话，使用 userId 作为键
    private static final Map<String, Session> sessions = new ConcurrentHashMap<>();

    @OnOpen
    public void onOpen(Session session) {
        // 获取 userId 参数并存储会话
        String userId = session.getQueryString().split("=")[1];
        sessions.put(userId, session);
        System.out.println("Connection opened, user ID: " + sessions);
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        System.out.println("Message received from client: " + message);
    }

    @OnClose
    public void onClose(Session session) {
        String userId = session.getQueryString().split("=")[1];
        sessions.remove(userId);
        System.out.println("Connection closed, user ID: " + userId);
    }

    @OnError
    public void onError(Session session, Throwable throwable) {
        System.out.println("Error occurred: " + throwable.getMessage());
    }

    // 向指定用户发送消息
    public static void sendMessageToUser(String userId, String message) throws IOException {
        Session session = sessions.get(userId);
        System.out.println(session);
        if (session != null && session.isOpen()) {
            session.getBasicRemote().sendText(message);
        }

    }
}
